import os
import joblib
import tensorflow as tf
# from onnxmltools.convert import convert_lightgbm
# from onnxmltools.convert.common.data_types import FloatTensorType
# from onnxruntime import InferenceSession


class ModelPersister:
    def __init__(self, save_path: str, file_name_prefix):
        self.save_path = save_path
        self.file_name_prefix = file_name_prefix

    def save_model(self, model):
        if isinstance(model, tf.keras.Model):
            self.save_model_keras(model)
        else:
            self.save_model_sklearn(model)

    def save_model_keras(self, model):
        file_path = os.path.join(self.save_path, f"model_{self.file_name_prefix}.keras")
        model.save(file_path)

    def load_model_keras(self):
        file_path = os.path.join(self.save_path, f"model_{self.file_name_prefix}.keras")
        model = tf.keras.models.load_model(file_path)
        return model

    def save_model_sklearn(self, model):
        file_path = os.path.join(self.save_path, f"model_{self.file_name_prefix}.pkl")
        joblib.dump(model, file_path)

    def load_model_sklearn(self):
        file_path = os.path.join(self.save_path, f"model_{self.file_name_prefix}.pkl")
        model = joblib.load(file_path)
        return model

    def save_scaler(self, scaler):
        file_path = os.path.join(self.save_path, f"scaler_{self.file_name_prefix}.pkl")
        joblib.dump(scaler, file_path)

    def load_scaler(self):
        file_path = os.path.join(self.save_path, f"scaler_{self.file_name_prefix}.pkl")
        return joblib.load(file_path)

    def save_label_encoder(self, encoder):
        file_path = os.path.join(self.save_path, f"encoder_{self.file_name_prefix}.pkl")
        joblib.dump(encoder, file_path)

    def load_label_encoder(self):
        file_path = os.path.join(self.save_path, f"encoder_{self.file_name_prefix}.pkl")
        return joblib.load(file_path)
